# -*- coding: utf-8 -*-
"""
Created on Mon Dec 25 13:56:27 2023

@author: HD
"""

import numpy as np
from statsmodels.tsa.arima.model import ARIMA


def arima_forecast(data, order=(1, 1, 1), forecast_steps=3):
    """
    使用ARIMA模型对多个时间序列进行预测。

    参数：
    - data: 输入的时间序列数据，维度为（n_timesteps, n_samples）。
    - order: ARIMA模型的阶数，形如 (p, d, q)。
    - forecast_steps: 预测的时间步数量。

    返回：
    - 预测结果，维度为（forecast_steps，n_samples）。
    """

    # 存储预测结果
    predictions = []

    # 对每个时间序列进行训练和预测
    for i in range(data.shape[1]):
        # 训练ARIMA模型
        model = ARIMA(data[:, i], order=order)
        model_fit = model.fit()

        # 进行未来预测
        forecast = model_fit.get_forecast(steps=forecast_steps)

        # 获取预测结果
        forecast_mean = forecast.predicted_mean
        predictions.append(forecast_mean)

    # 将预测结果转换为数组，维度为（forecast_steps，n_samples）
    predictions = np.array(predictions).T
    return predictions


def forecast_with_autoarima(data, s, d_range, D_range, p_range, q_range, P_range, Q_range, forecast_steps):
    """
    使用AutoARIMA模型进行预测

    参数：
    - df: 输入的DataFrame，包含时间序列数据，形状为(N, F)，N代表时间步，F代表不同的时间序列
    - s: 季节性周期长度
    - d_range: 差分阶数的范围，例如(0, 1)表示可选0或1阶差分
    - D_range: 季节性差分阶数的范围，例如(0, 1)表示可选0或1阶季节性差分
    - p_range: 自回归阶数的范围，例如(0, 2, 4)表示可选0、1或2阶自回归
    - q_range: 移动平均阶数的范围，例如(0, 2, 4)表示可选0、1或2阶移动平均
    - P_range: 季节性自回归阶数的范围，例如range(2)表示可选0或1阶季节性自回归
    - Q_range: 季节性移动平均阶数的范围，例如range(2)表示可选0或1阶季节性移动平均
    - forecast_steps: 预测的时间步数

    返回值：
    - fc: 预测结果，形状为(F, forecast_steps)，每列包含对应时间序列的未来预测值
    """
    from cuml.tsa.auto_arima import AutoARIMA

    model = AutoARIMA(data)
    # 搜索最佳模型参数
    model.search(s=s, d=d_range, D=D_range, p=p_range, q=q_range, P=P_range, Q=Q_range, method="css", truncate=20)
    # 拟合模型
    model.fit(method="css-ml")
    # 进行未来预测
    fc = model.forecast(forecast_steps)
    return fc


def forecast(data, forecast_steps=3, use_gpu=False):
    """
    对多个时间序列进行预测，可选择使用ARIMA模型或AutoARIMA模型。

    参数：
    - data: 输入的时间序列数据，维度为（n_timesteps, n_samples）。
    - order: ARIMA模型的阶数，形如 (p, d, q)。
    - forecast_steps: 预测的时间步数量。
    - use_gpu: 是否使用AutoARIMA模型。

    返回：
    - 预测结果，维度为（forecast_steps，n_samples）。
    """
    if use_gpu:
        return forecast_with_autoarima(data, s=2, d_range=(0, 5), D_range=(0, 5), p_range=(0, 5), q_range=(0, 5),
                                       P_range=(0, 5), Q_range=(0, 5), forecast_steps=forecast_steps)
    else:
        order = (1, 0, 0)  # 无GPU情况下可以调整的参数,备选（1,0,0）（0,1,1）
        return arima_forecast(data, order=order, forecast_steps=forecast_steps)
